# whisper.load_model方法会自动下载相应的模型,
# 可选：['tiny.en', 'tiny', 'base.en', 'base', 'small.en', 'small', 'medium.en', 'medium', 'large-v1', 'large-v2', 'large-v3', 'large', 'large-v3-turbo', 'turbo']。
import whisper
import json
import wave
from vosk import Model, KaldiRecognizer

# Whisper由OpenAI开发,其训练数据的多样性使得它在处理各种口音、背景噪音和专业术语方面具有较好的鲁棒性。
# large-v3-turbo是Whisper模型的一个改进版本,在保持较高准确性的同时,速度有了很大的提升,这使得它在实时性要求较高的场景下非常有优势。
# 如果需要快速准确地处理语音,并且有足够的计算资源支持大型模型,那么large-v3-turbo是很好的选择。
#
# Vosk是一个轻量级的语音识别工具包,基于Kaldi集成,在处理一些复杂语言场景或特殊口音时,可能相对Whisper会稍显不足。


def split_chinese(model_name='base'):
    model = whisper.load_model(model_name)
    result = model.transcribe("data/test_voice.mp3", fp16=False)
    print(result['text'])
    print(result['language'])
    # 可通过opencc库将繁体字转换为简体字：
    # pip install opencc-python-reimplemented
    # from opencc import OpenCC
    #
    # # 初始化转换器,从繁体中文转换到简体中文
    # converter = OpenCC('t2s')
    # model = whisper.load_model("base")
    # result = model.transcribe("test.wav")
    # print(converter.convert(result['text']))


def recognize_wave(model, file_path):
    # 打开WAV音频文件
    with wave.open(file_path, 'rb') as wf:
        rate = wf.getframerate()  # 采样率
        frames = wf.readframes(wf.getnframes())  # 读取所有帧

    rec = KaldiRecognizer(model, rate)
    rec.SetWords(True)
    str_ret = ""

    # 识别音频数据
    if rec.AcceptWaveform(frames):
        result = json.loads(rec.Result())
        if 'text' in result:
            str_ret += result['text']

    result = json.loads(rec.FinalResult())
    if 'text' in result:
        str_ret += result['text']

    str_ret = "".join(str_ret.split())
    return str_ret


if __name__ == '__main__':
    # large-v3-turbo模型比base模型的识别准确性更高,并且识别结果为简体中文,不用额外进行繁转简,更加友好的是,识别结果有标点符号,可读性更强。
    split_chinese("base")
    split_chinese("large-v1")

    model = Model("vosk-model-cn-0.22")
    file_path = 'test.wav'  # 请确保文件名和路径正确
    res = recognize_wave(model, file_path)
    print(res)
